#ifndef __PE_TOOLS__
#define __PE_TOOLS__

#include <algorithm>
#include <deque>
#include <set>
#include <vector>
#include "../../comm/xcrc32.h"

namespace xm
{
	typedef std::set<UINT32>		UINT32SET;
	typedef std::deque<UINT32>		UINT32STACK;

	class WinPETools : public IKPETools, UnkBase
	{
	public:
		
		typedef IPEDisAsmCallback		TCallback;

	protected:

		typedef std::vector<IMPITEM>	imports_t;
		typedef std::vector<IMPNAME>	imp_names_t;

		iptr<IPeReader>		_archpe;
		imports_t			_imports;
		imp_names_t			_impnames;
		uint32_t			_do_mask;
		//////////////////////////////////////////////////////////////////////////
		//
		//	image base & size
		//
		IPeReader::QWORD	_imgbase;
		DWORD				_imgsize;
		//////////////////////////////////////////////////////////////////////////
		//
		//	diasm
		//
		uint32_t			_dasm_opts;
		UINT32STACK			_stack;
		UINT32SET			_traced;

		static bool _less_by_thunk( const IMPITEM & l, const IMPITEM & r )
		{
			return l.iat < r.iat;
		}
		static bool _less_by_namrva( const IMPNAME & l, const IMPNAME & r )
		{
			return l.namrva < r.namrva;
		}

		enum 
		{
			_MSK_FETCH_IMPORTS = 1,
		};

	public:

		UNKNOWN_IMP1(IKPETools);
		DEFAULT_INIT_CLASS();

	public:

		WinPETools() : _do_mask(0), _imgbase(0), _imgsize(0), _dasm_opts(0)
		{
			
		}

		STDMETHOD(Init)( IPeReader * ope )
		{
			Close();
			_archpe = ope;
			if( !_archpe ) return E_NOINTERFACE;
			if( !_archpe->ImageBaseVa(&_imgbase,&_imgsize) )
				return E_FAIL;
			return S_OK;
		}

		STDMETHOD_(IPeReader *, WinPeObject)()
		{
			return _archpe;
		}

		STDMETHOD_(SIZE_T, ImpGetCount)()
		{
			_load_imports();
			return _imports.size();
		}
		STDMETHOD_(IMPITEM*, ImpGetAt)( SIZE_T idx = 0 )
		{
			_load_imports();
			if( idx >= _imports.size() )
				return NULL;
			return &_imports[idx];
		}
		STDMETHOD_(IMPITEM*, ImpFindByIAT)( UINT64 addr, BOOL is_va )
		{
			_load_imports();

			if( is_va ) 
			{
				if( addr < _imgbase ) 
					return NULL;
				addr -= _imgbase;
			}
			
			if( addr >= _imgsize ) 
				return NULL;
			
			UINT32 iat = (UINT32)addr;
			IMPITEM key = { iat };
			imports_t::iterator it = std::lower_bound( _imports.begin(), _imports.end(), key, _less_by_thunk );
			if( it == _imports.end() ) return NULL;
			if( it->iat != iat ) return NULL;
			return &(*it);
		}
		STDMETHOD_(UINT32, ImpFindByName)( UINT32 namerva )
		{
			_load_imports();
			IMPNAME key = { namerva };
			imp_names_t::iterator it = std::lower_bound( _impnames.begin(), _impnames.end(), key, _less_by_namrva );
			if( it == _impnames.end() ) return NULL;
			if( it->namrva != namerva ) return NULL;
			return it->iat;
		}
		STDMETHOD_(VOID,Close)()
		{
			_archpe.dispose();
			_imports.clear();
			_impnames.clear();
			_do_mask = 0;
			_dasm_opts = 0;
			_imgbase = 0;
			_imgsize = 0;
			_traced.clear();
			_stack.clear();
		}

		STDMETHOD(FetchBIL)( UINT32 rva, SIZE_T options, IPEDisAsmCallback * callback );

		STDMETHOD_(LPCSTR, IsMoudleImported)( LPCSTR aName, MATCH_MODE_T mode )
		{
			LPCSTR aMod = NULL;
			while( aMod = _archpe->NextImportModule( aMod ) )
			{
				if( mode == KeyStrEqulModuleName )
				{
					if( !stricmp( aName, aMod ) )
						return aMod;
				}
				else if( mode == KeyStrInModuleName )
				{
					if( tcs::istr( aMod, aName ) )
						return aMod;
				}
				else if( mode == ModuleNameInKeyStr )
				{
					if( tcs::istr( aName, aMod ) )
						return aMod;
				}
			}
			return NULL;
		}

	protected:

		void _load_imports( )
		{
			if( _do_mask & _MSK_FETCH_IMPORTS ) return ;
			_do_mask |= _MSK_FETCH_IMPORTS;

			DWORD nimp = 0;
			RASSERTV( _archpe->ImportCount( &nimp ) );
			_imports.reserve( nimp ); 
			_impnames.reserve( nimp );

			LPCSTR aMod = NULL;
			while( aMod = _archpe->NextImportModule( aMod ) )
			{
				LPCSTR aApi = NULL;
				while( aApi = _archpe->NextImportFunction(aMod, aApi) )
				{
					DWORD iat = 0, namrva = 0;
					if( _archpe->ImportFunctionSlotRva(aMod, aApi, &iat ) )
					{
						UINT32 crc = gnu_crc32( aApi, strlen(aApi), 'APIC' );
						IMPITEM item = { iat, aApi, aMod, crc };
						_imports.push_back( item );
						if( _archpe->ImportFunctionNameRVA(aMod, aApi, &namrva ) )
						{
							IMPNAME name = { namrva, iat };
							_impnames.push_back( name );
						}
					}
				}
			}

			std::sort( _imports.begin(), _imports.end(), _less_by_thunk );
			std::sort( _impnames.begin(), _impnames.end(), _less_by_namrva );
		}

	protected:

		HRESULT Fetch( UINT32 uRVA, TCallback* callback );
		HRESULT InsertInstruction32( UINT32 uRVA, hde32s & hid, UINT32& uOpcode, TCallback* lpCallback );
		HRESULT InsertInstruction64( UINT32 uRVA, hde64s & hid, UINT32& uOpcode, TCallback* lpCallback );
		VOID	TraceRVA( UINT32 uRVA );
		HRESULT DisStackTop(TCallback* pcb);
		BOOL	ProcessBranch32( UINT32& uRVA, hde32s& hid, UINT32 uOpcode );
		BOOL	ProcessBranch64( UINT32& uRVA, hde64s& hid, UINT32 uOpcode );

	};



};


#endif